"""
Agent 后端服务 - 基于 FastAPI 提供 API 接口（简化版）

提供同步处理模式的代理服务，直接返回结果或中断数据，
简化了状态管理和通信流程
"""
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from typing import Dict, Any, Optional, List, Literal
import json
import uuid
import time

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import interrupt, Command
from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
from langchain_tavily import TavilySearch
import os
import uvicorn

# 会话存储 - 保存每个用户的智能体实例和状态
# 暂时只支持一个用户一个会话，不能一个用户多个会话
# 结构: {user_id: {
#   "agent": agent,  
#   "session_id": session_id, 
#   "status": "idle|running|interrupted|completed|error",
#   "last_response": AgentResponse,     # 上次响应
#   "last_query": str,                  # 上次查询
#   "last_updated": timestamp           # 上次更新时间
# }}
sessions = {}

# 创建 LLM
model = ChatOpenAI(model="gpt-4o-mini")

# 定义生命周期上下文管理器
from contextlib import asynccontextmanager
from langgraph.checkpoint.postgres import PostgresSaver

@asynccontextmanager
async def lifespan(app: FastAPI):
    
    # 初始化 PostgreSQL 检查点保存器
    with PostgresSaver.from_conn_string("postgresql://postgres:yourpassword@localhost/postgres?sslmode=disable") as checkpointer:
        checkpointer.setup()

        # 将检查点保存器设置为全局可访问
        app.state.checkpointer = checkpointer

        yield

# 使用生命周期上下文管理器初始化应用
app = FastAPI(
    title="Agent 后端服务", 
    description="基于 LangGraph 提供 AI Agent 服务",
    lifespan=lifespan
)

# 定义请求和响应模型
class AgentRequest(BaseModel):
    user_id: str
    query: str
    system_message: Optional[str] = "你会使用工具来帮助用户。如果工具使用被拒绝，请提示用户。"

class AgentResponse(BaseModel):
    session_id: str
    status: str  # interrupted, completed, error, running, idle
    timestamp: float = Field(default_factory=lambda: time.time())

    #error
    message: Optional[str] = None

    #completed
    result: Optional[Dict[str, Any]] = None

    #interrupted
    interrupt_data: Optional[Dict[str, Any]] = None
    
    class Config:
        json_schema_extra = {
            "example": {
                "session_id": "123e4567-e89b-12d3-a456-426614174000",
                "status": "interrupted",
                "timestamp": 1621500000.123,
                "interrupt_data": {
                    "interrupt_type": "tool_use_approval",
                    "tool": "tavily_search",
                    "args": {"query": "哪吒3最新信息", "search_depth": "basic"},
                    "message": "准备使用Tavily搜索...",
                    "expected_responses": ["accept", "reject", "edit"]
                }
            }
        }

class InterruptResponse(BaseModel):
    user_id: str
    session_id: str
    response_type: str  # accept, reject, edit
    args: Optional[Dict[str, Any]] = None

class SystemInfoResponse(BaseModel):
    sessions_count: int
    active_users: List[str]
    
# 处理智能体结果的辅助函数
def process_agent_result(
    session_id: str, 
    result: Dict[str, Any],
    user_id: Optional[str] = None
) -> AgentResponse:
    """
    处理智能体执行结果，统一处理中断和结果
    
    Args:
        session_id: 会话ID
        result: 智能体执行结果
        user_id: 用户ID，如果提供，将更新会话状态
        
    Returns:
        AgentResponse: 标准化的响应对象
    """
    response = None

    try:
        # 检查是否有中断
        if "__interrupt__" in result:
            interrupt_data = result["__interrupt__"][0].value
            
            # 确保中断数据有类型信息
            if "interrupt_type" not in interrupt_data:
                interrupt_data["interrupt_type"] = "unknown"
            
            # 返回中断信息
            response = AgentResponse(
                session_id=session_id,
                status="interrupted",
                interrupt_data=interrupt_data
            )
        else:
            # 如果没有中断，返回最终结果
            response = AgentResponse(
                session_id=session_id,
                status="completed",
                result=result
            )
        
    except Exception as e:
        response = AgentResponse(
            session_id=session_id,
            status="error",
            message=f"处理智能体结果时出错: {str(e)}"
        )
        
    # 如果提供了用户ID，更新会话状态
    if user_id and user_id in sessions:
        sessions[user_id]["status"] = response.status
        sessions[user_id]["last_response"] = response
        sessions[user_id]["last_updated"] = time.time()
        
    return response

# Tavily搜索工具，需要人工审核/批准
def tavily_search(query: str, search_depth: Optional[str] = "basic"):
    """
    使用Tavily进行网络搜索
    
    参数:
    - query: 搜索查询
    - search_depth: 搜索深度，可选值为 "basic"（基础搜索）或 "advanced"（深度搜索）
    """
    # 中断执行，等待人工审核
    response = interrupt({
        "interrupt_type": "tool_use_approval",
        "tool": "tavily_search",
        "args": {
            "query": query,
            "search_depth": search_depth
        },
        "message": f"准备使用Tavily搜索：\n- 查询内容: {query}\n- 搜索深度: {search_depth}\n\n是否允许继续？\n输入 'yes' 接受，'no' 拒绝，或 'edit' 修改查询关键词",
        "expected_responses": ["accept", "reject", "edit"]
    })   

    # 处理人工响应
    if response["type"] == "accept":
        pass
    elif response["type"] == "edit":
        query = response["args"]["query"]
    else:
        return f"该工具被拒绝使用，请尝试其他方法或拒绝回答问题。"
    
    # 使用LangChain的TavilySearch工具进行实际搜索
    try:
        search = TavilySearch(api_key=os.environ.get("TAVILY_API_KEY"))
        search_results = search.invoke({
            "query": query,
            "search_depth": search_depth,
            "max_results": 5
        })
    except Exception as e:
       return f"该工具调用失败: {str(e)}"
    
    return json.dumps(search_results, ensure_ascii=False, indent=2)

def create_tavily_search_agent(user_id: str):
    """
    创建一个带有Tavily搜索工具的智能体
    
    Args:
        user_id: 用户ID，用于隔离不同用户的会话
        
    Returns:
        agent: 已配置的智能体实例
    """
    # 创建检查点保存器
    #checkpointer = InMemorySaver() 

    # 创建智能体
    agent = create_react_agent(
        model=model,
        tools=[tavily_search],
        checkpointer=app.state.checkpointer,  # 使用全局检查点保存器
    )
    
    return agent

@app.post("/agent/invoke", response_model=AgentResponse)
def invoke_agent(request: AgentRequest):
    """启动智能体处理用户请求 - 同步版本，等待执行完成或中断"""
    user_id = request.user_id
    
    # 检查是否已有会话，如果没有则创建新会话
    if user_id not in sessions:
        # 只在创建新会话时生成新的会话ID
        session_id = str(uuid.uuid4())
        agent = create_tavily_search_agent(user_id)
        sessions[user_id] = {
            "agent": agent,
            "session_id": session_id,
            "status": "idle",
            "last_query": None,
            "last_response": None,
            "last_updated": time.time()
        }
    else:
        # 使用现有会话的ID和agent
        agent = sessions[user_id]["agent"]
        session_id = sessions[user_id]["session_id"]
    
    # 新请求统一更新会话信息
    sessions[user_id]["status"] = "running"
    sessions[user_id]["last_query"] = request.query
    sessions[user_id]["last_response"] = None
    sessions[user_id]["last_updated"] = time.time()
    
    # 初始化智能体输入
    messages = [
        {"role": "system", "content": request.system_message},
        {"role": "user", "content": request.query}
    ]
    
    try:
        # 先调用智能体
        result = agent.invoke({"messages": messages}, config={"configurable": {"thread_id": session_id}})

        # 再处理结果并更新会话状态
        return process_agent_result(session_id, result, user_id)
    
    except Exception as e:
        # 异常处理
        error_response = AgentResponse(
            session_id=session_id,
            status="error",
            message=f"处理请求时出错: {str(e)}"
        )
        
        # 更新会话状态
        sessions[user_id]["status"] = "error"
        sessions[user_id]["last_response"] = error_response
        sessions[user_id]["last_updated"] = time.time()
        
        return error_response

@app.post("/agent/resume", response_model=AgentResponse)
def resume_agent(response: InterruptResponse):
    """恢复被中断的智能体执行 - 同步版本，等待执行完成或再次中断"""
    user_id = response.user_id
    client_session_id = response.session_id
    
    # 检查用户会话是否存在
    if user_id not in sessions:
        raise HTTPException(status_code=404, detail=f"用户会话 {user_id} 不存在")
    
    # 检查会话ID是否匹配
    server_session_id = sessions[user_id]["session_id"]
    if server_session_id != client_session_id:
        raise HTTPException(status_code=400, detail="会话ID不匹配，可能是过期的请求")
    
    # 检查会话状态是否为中断
    if sessions[user_id]["status"] != "interrupted":
        raise HTTPException(status_code=400, detail=f"会话当前状态为 {sessions[user_id]['status']}，无法恢复非中断状态的会话")
    
    # 更新会话状态为运行中
    sessions[user_id]["status"] = "running"
    sessions[user_id]["last_updated"] = time.time()
    
    # 获取智能体和配置
    agent = sessions[user_id]["agent"]
    
    # 构造响应数据
    command_data = {
        "type": response.response_type
    }
    
    # 如果提供了参数，添加到响应数据中
    if response.args:
        command_data["args"] = response.args
    
    try:
        # 先恢复智能体执行
        result = agent.invoke(Command(resume=command_data), config={"configurable": {"thread_id": server_session_id}})

        # 再处理结果并更新会话状态
        return process_agent_result(server_session_id, result, user_id)
    
    except Exception as e:
        # 异常处理
        error_response = AgentResponse(
            session_id=server_session_id,
            status="error",
            message=f"恢复执行时出错: {str(e)}"
        )
        
        # 更新会话状态
        sessions[user_id]["status"] = "error"
        sessions[user_id]["last_response"] = error_response
        sessions[user_id]["last_updated"] = time.time()
        
        return error_response

class SessionStatusResponse(BaseModel):
    user_id: str
    session_id: Optional[str] = None
    status: str  # not_found, idle, running, interrupted, completed, error
    message: Optional[str] = None
    last_query: Optional[str] = None
    last_updated: Optional[float] = None
    last_response: Optional[AgentResponse] = None

@app.get("/agent/status/{user_id}", response_model=SessionStatusResponse)
def get_agent_status(user_id: str):
    """获取用户智能体的会话状态"""
    if user_id not in sessions:
        return SessionStatusResponse(
            user_id=user_id,
            status="not_found",
            message=f"用户 {user_id} 的会话不存在"
        )
    
    session = sessions[user_id]
    return SessionStatusResponse(
        user_id=user_id,
        session_id=session["session_id"],
        status=session["status"],
        last_query=session["last_query"],
        last_updated=session["last_updated"],
        last_response=session["last_response"]
    )

@app.get("/system/info", response_model=SystemInfoResponse)
def get_system_info():
    """获取系统状态信息"""
    return SystemInfoResponse(
        sessions_count=len(sessions),
        active_users=list(sessions.keys())
    )

@app.delete("/agent/session/{user_id}")
def delete_agent_session(user_id: str):
    """删除用户会话"""
    if user_id not in sessions:
        raise HTTPException(status_code=404, detail=f"用户 {user_id} 的会话不存在")
    
    # 删除会话
    del sessions[user_id]
    
    return {
        "status": "success",
        "message": f"用户 {user_id} 的会话已删除"
    }

# 启动服务器
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8001)
